use std::num::NonZeroU64;

use super::arc::{ArcBindGroup, ArcBindGroupLayout, ArcBuffer, ArcSampler, ArcTextureView};

pub trait AsGpuBindGroup {
    fn as_group(&self) -> &ArcBindGroup;
    fn as_group_layout(&self) -> &ArcBindGroupLayout;
}

/// Builder pattern for bind group layouts; basically just produces a Vec<BindGroupLayoutEntry>.
#[derive(Debug)]
pub struct BindGroupLayoutBuilder {
    entries: Vec<wgpu::BindGroupLayoutEntry>,
}

impl BindGroupLayoutBuilder {
    pub fn new() -> Self {
        BindGroupLayoutBuilder { entries: vec![] }
    }

    pub fn buffer(
        mut self,
        visibility: wgpu::ShaderStages,
        ty: wgpu::BufferBindingType,
        has_dynamic_offset: bool,
    ) -> Self {
        self.entries.push(wgpu::BindGroupLayoutEntry {
            binding: self.entries.len() as _,
            visibility,
            ty: wgpu::BindingType::Buffer {
                ty,
                has_dynamic_offset,
                min_binding_size: None,
            },
            count: None,
        });
        self
    }

    pub fn image(
        mut self,
        visibility: wgpu::ShaderStages,
        view_dimension: wgpu::TextureViewDimension,
    ) -> Self {
        self.entries.push(wgpu::BindGroupLayoutEntry {
            binding: self.entries.len() as _,
            visibility,
            ty: wgpu::BindingType::Texture {
                sample_type: wgpu::TextureSampleType::Float { filterable: true },
                view_dimension,
                multisampled: false,
            },
            count: None,
        });
        self
    }

    pub fn sampler(mut self, visibility: wgpu::ShaderStages) -> Self {
        self.entries.push(wgpu::BindGroupLayoutEntry {
            binding: self.entries.len() as _,
            visibility,
            ty: wgpu::BindingType::Sampler(wgpu::SamplerBindingType::Filtering),
            count: None,
        });
        self
    }

    pub fn create(self, device: &wgpu::Device) -> ArcBindGroupLayout {
        ArcBindGroupLayout::new(
            device.create_bind_group_layout(&wgpu::BindGroupLayoutDescriptor {
                label: None,
                entries: &self.entries,
            }),
        )
    }
}

/// This is used as a key into the HashMap cache to uniquely identify a bind group.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum BindGroupEntryKey {
    Buffer {
        id: u64,
        offset: u64,
        size: Option<u64>,
    },
    Image {
        id: u64,
    },
    Sampler {
        id: u64,
    },
}
#[derive(Debug)]
pub struct BindGroupBuilder<'a> {
    layout: BindGroupLayoutBuilder,
    entries: Vec<wgpu::BindGroupEntry<'a>>,
    key: Vec<BindGroupEntryKey>,
}

impl<'a> BindGroupBuilder<'a> {
    pub fn new() -> Self {
        BindGroupBuilder {
            layout: BindGroupLayoutBuilder::new(),
            entries: vec![],
            key: vec![],
        }
    }

    pub fn buffer(
        mut self,
        buffer: &'a ArcBuffer,
        offset: u64,
        visibility: wgpu::ShaderStages,
        ty: wgpu::BufferBindingType,
        has_dynamic_offset: bool,
        size: Option<u64>,
    ) -> Self {
        self.entries.push(wgpu::BindGroupEntry {
            binding: self.entries.len() as _,
            resource: wgpu::BindingResource::Buffer(wgpu::BufferBinding {
                buffer: buffer.as_ref(),
                offset,
                size: size.map(|x| NonZeroU64::new(x).unwrap()),
            }),
        });

        self.key.push(BindGroupEntryKey::Buffer {
            id: buffer.id(),
            offset,
            size,
        });

        BindGroupBuilder {
            layout: self.layout.buffer(visibility, ty, has_dynamic_offset),
            entries: self.entries,
            key: self.key,
        }
    }

    pub fn image(
        mut self,
        view: &'a ArcTextureView,
        visibility: wgpu::ShaderStages,
        view_dimension: wgpu::TextureViewDimension,
    ) -> Self {
        self.entries.push(wgpu::BindGroupEntry {
            binding: self.entries.len() as _,
            resource: wgpu::BindingResource::TextureView(view.as_ref()),
        });

        self.key.push(BindGroupEntryKey::Image { id: view.id() });

        BindGroupBuilder {
            layout: self.layout.image(visibility, view_dimension),
            entries: self.entries,
            key: self.key,
        }
    }

    pub fn sampler(mut self, sampler: &'a ArcSampler, visibility: wgpu::ShaderStages) -> Self {
        self.entries.push(wgpu::BindGroupEntry {
            binding: self.entries.len() as _,
            resource: wgpu::BindingResource::Sampler(sampler.as_ref()),
        });

        self.key
            .push(BindGroupEntryKey::Sampler { id: sampler.id() });

        BindGroupBuilder {
            layout: self.layout.sampler(visibility),
            entries: self.entries,
            key: self.key,
        }
    }

    pub fn create(self, device: &wgpu::Device) -> (ArcBindGroup, ArcBindGroupLayout) {
        let layout = self.layout.create(device);
        let group = ArcBindGroup::new(device.create_bind_group(&wgpu::BindGroupDescriptor {
            label: None,
            layout: layout.as_ref(),
            entries: &self.entries,
        }));
        (group, layout)
    }

    #[inline]
    pub fn entries(&self) -> &[wgpu::BindGroupEntry] {
        &self.entries
    }
}
